package main

import (
	"bytes"
	"encoding/binary"
	"errors"
	"github.com/pterm/pterm"
	"io"
	"time"
	"xmediaEmu/pkg/decoder/codec"
	"xmediaEmu/pkg/decoder/codec/aacparser"
	"xmediaEmu/pkg/decoder/codec/h264parser"
	"xmediaEmu/pkg/format/rtmp"
	"xmediaEmu/pkg/media/h264writer"
)

type Options struct {
	OutputCodecTypes []codec.CodecType
}

func main() {
	// duration := time.Duration(0) // packet time ,比如25fps的为40ms.
	// 结果中提取h.264.
	// 需要音频文件的话需要另外提取.
	outputfile := "rtmp30s.h264"
	h264File, err := h264writer.New(outputfile)
	if err != nil {
		pterm.FgRed.Println(err)
		panic(err)
	}
	defer h264File.Close()

	// options := Options{}
	url := "rtmp://wsyidong.grtn.cn/live/gdws" // live=1

	// 只负责拉取.
	demuxer, err := rtmp.Dial(url)
	if err != nil {
		pterm.FgRed.Println(err)
		panic(err)
	}
	defer demuxer.Close()
	//err = demuxer.Prepare()
	//if err != nil {
	//	pterm.FgRed.Println(err)
	//	panic(err)
	//}
	//
	//err = demuxer.Prepare()
	//if err != nil {
	//	pterm.FgRed.Println(err)
	//	panic(err)
	//}

	// video 打印
	pterm.FgYellow.Println("Streams get:")
	var streams []codec.CodecData
	if streams, err = demuxer.Streams(); err != nil && err != io.EOF {
		pterm.FgRed.Println(err)
		return
	}
	pterm.FgYellow.Println("Streams:", len(streams))

	if err = rtmpStreamer.WriteHeader(streams); err != nil {
		pterm.FgRed.Println(err)
		return
	}

	// 打印标签.
	for _, stream := range streams {
		// fmt.Print(stream.Type(), " ")
		if stream.Type() == codec.H264 {
			pterm.FgYellow.Println("found dest stream type:", stream.Type(), " ")
		} else {
			pterm.FgWhite.Println(stream.Type(), " ")
		}
	}
	// video 打印 end

	// codec 写死为rtmp的编解码flv.
	// options.OutputCodecTypes = rtmp.CodecTypes
	// 循环读取30s视频.
	var currentBytes int
	ticker := time.NewTicker(time.Second * 30)
	for {
		select {
		case <-ticker.C:
			// 30 second 到了停止.
			pterm.FgWhite.Printf("Receive the message success, the total lenth is %d\n", currentBytes)
			pterm.FgYellow.Println("ticks stop!")
			return
		default:
			var pkt codec.Packet
			// pterm.FgYellow.Println("ReadPacket:")
			if pkt, err = demuxer.ReadPacket(); err != nil {
				if err == io.EOF {
					err = nil
					pterm.FgRed.Println("end EOF")
					break
				}
				pterm.FgRed.Println(err)
				return
			}

			streamId := int(pkt.Idx)

			// 打印当前编码，非视频则跳过.
			if demuxer.VideoStreamId() != streamId {
				pterm.FgWhite.Printfln("stream[%d] id not video: VideoStreamId[%d]", streamId, demuxer.VideoStreamId())
				continue
			}

			// 以下代码解析关键帧会有许多问题.
			//if pkt.IsKeyFrame {
			//	sps, pps, err := rtmpStreamer.GetSpsPpsFromPacket(pkt)
			//	if err != nil {
			//		pterm.FgRed.Println(err)
			//		continue
			//	}
			//	if err = h264File.WriteFile(sps); err != nil {
			//		pterm.FgLightRed.Printf("hWriter.Write: Write Sps failed:%v\n. ", err)
			//		continue
			//	} else {
			//		pterm.FgGreen.Printf("hWriter.Write Sps count:%d done. \n", len(sps))
			//	}
			//
			//	if err = h264File.WriteFile(pps); err != nil {
			//		pterm.FgLightRed.Printf("hWriter.Write: Write pps failed:%v\n. ", err)
			//	} else {
			//		pterm.FgGreen.Printf("hWriter.Write pps count:%d done. \n", len(pps))
			//	}
			//	continue
			//}

			var payloads [][]byte
			payloads = rtmpStreamer.GetH264FromPacket(pkt)
			pterm.FgYellow.Println("packet information: ", pkt.Idx, pkt.Time, len(pkt.Data), pkt.IsKeyFrame, len(payloads))
			if len(payloads) == 0 {
				continue
			}

			// 提取的如果是视频h264直接写入文件.
			// tag.data是否就是h264原始文件?.
			// 既不是原始h264文件也不是flv文件，是啥格式？..
			for _, videoElement := range payloads {
				if err = h264File.WriteFile(videoElement); err != nil {
					pterm.FgLightRed.Printf("receiveLocalMessage: WriteRTP failed:%v\n. ", err)
				} else {
					pterm.FgYellow.Printf("hWriter.WriteFile Payload count:%d done. \n", len(pkt.Data))
				}
			}
			currentBytes += len(pkt.Data)
		}
	}

	// h264文件结尾？.

	pterm.FgYellow.Println("Normal exit!")
}

// 写死编解码信息.
var rtmpStreamer = NewRtmpStreamer(Capabilities["audio"], Capabilities["video"])
var Capabilities = map[string]*Capability{
	"audio": &Capability{
		Codecs: []string{"AAC"},
		Rtcpfbs: []*RtcpFeedback{
			&RtcpFeedback{
				ID: "nack",
			},
		},
	},
	"video": &Capability{
		Codecs: []string{"H264"},
		Rtx:    true,
		Rtcpfbs: []*RtcpFeedback{
			&RtcpFeedback{
				ID: "transport-cc",
			},
			&RtcpFeedback{
				ID:     "ccm",
				Params: []string{"fir"},
			},
			&RtcpFeedback{
				ID: "nack",
			},
		},
		Extensions: []string{
			"urn:3gpp:video-orientation",
			"http://www.ietf.org/id/draft-holmer-rmcat-transport-wide-cc-extensions-01",
			"http://www.webrtc.org/experiments/rtp-hdrext/abs-send-time",
		},
	},
}

// RtmpStream 定义rtmp解析的结构streamer.
type RtmpStream struct {
	streams        []codec.CodecData
	videoCodecData h264parser.CodecData
	audioCodecData aacparser.CodecData

	spspps bool

	// aac头部
	adtsheader []byte

	// TODO: 以下用于构建SDP。不需要则跳过.
	// 后续和外部媒体通信需要根据外部媒体调节自己的编码属性，尽量不做转码，比如h265.
	audioCapability *Capability
	videoCapability *Capability
}

type RtcpFeedback struct {
	ID     string   `json:"id,omitempty"`
	Params []string `json:"params,omitempty"`
}

type Capability struct {
	Codecs     []string        `json:"codecs"`
	Rtx        bool            `json:"rtx,omitempty"`
	Rtcpfbs    []*RtcpFeedback `json:"rtcpfbs,omitempty"`
	Extensions []string        `json:"extensions,omitempty"`
	Simulcast  bool            `json:"simulcast,omitempty"`
}

// NewMediaTransform  create media transform
func NewRtmpStreamer(audio *Capability, video *Capability) *RtmpStream {
	streamer := &RtmpStream{}
	streamer.audioCapability = audio
	streamer.videoCapability = video
	return streamer
}

// WriteHeader got sps and pps
func (self *RtmpStream) WriteHeader(streams []codec.CodecData) error {
	self.streams = streams

	for _, stream := range streams {
		if stream.Type() == codec.H264 {
			h264Codec := stream.(h264parser.CodecData)
			self.videoCodecData = h264Codec
		}
		if stream.Type() == codec.AAC {
			aacCodec := stream.(aacparser.CodecData)
			self.audioCodecData = aacCodec
			self.adtsheader = make([]byte, 7)
		}
	}

	return nil
}

// GetH264FromPacket 解析rtmp的数据块部分.
// 音频AAC情况下返回的byte数组只有一个，不分片.
func (self *RtmpStream) GetH264FromPacket(packet codec.Packet) (result [][]byte) {
	stream := self.streams[packet.Idx]
	// pterm.FgYellow.Println("GetH264FromPacket stream type：", stream.Type())
	if stream.Type() == codec.H264 {
		// pterm.FgYellow.Println("GetH264FromPacket 1")
		if !self.spspps {
			var b bytes.Buffer
			b.Write([]byte{0, 0, 0, 1})
			b.Write(self.videoCodecData.SPS())
			b.Write([]byte{0, 0, 0, 1})
			b.Write(self.videoCodecData.PPS())
			result = append(result, b.Bytes())
			self.spspps = true
			// pterm.FgYellow.Println("GetH264FromPacket 2")
		}

		pktnalus, _ := h264parser.SplitNALUs(packet.Data)
		for _, nalu := range pktnalus {
			var b bytes.Buffer
			b.Write([]byte{0, 0, 0, 1})
			b.Write(nalu)
			result = append(result, b.Bytes())
			// pterm.FgYellow.Println("GetH264FromPacket 3")
		}

	}

	if stream.Type() == codec.AAC {
		var b bytes.Buffer
		adtsbuffer := []byte{}
		aacparser.FillADTSHeader(self.adtsheader, self.audioCodecData.Config, 1024, len(packet.Data))
		adtsbuffer = append(adtsbuffer, self.adtsheader...)
		adtsbuffer = append(adtsbuffer, packet.Data...)
		result = append(result, b.Bytes())
	}
	return
}

// 增加关键帧解析,专门针对视频.
func (self *RtmpStream) GetSpsPpsFromPacket(packet codec.Packet) (sps, pps []byte, err error) {
	stream := self.streams[packet.Idx]
	// pterm.FgYellow.Println("GetH264FromPacket stream type：", stream.Type())
	if stream.Type() == codec.H264 { // h265的后续加上.
		sps, pps, err = ParseSpsPpsFromSeqHeaderWithoutMalloc(packet.Data)

		// pterm.FgYellow.Println("GetH264FromPacket 1")
		//if !self.spspps {
		//	self.spspps = true
		//	// pterm.FgYellow.Println("GetH264FromPacket 2")
		//}
	} else {
		err = errors.New("unknown codec")
	}
	return
}

// ParseSpsPpsFromSeqHeaderWithoutMalloc
//
// 从AVCC格式的Seq Header中得到SPS和PPS内存块
//
// @param payload: rtmp message的payload部分或者flv tag的payload部分
//
//	注意，包含了头部2字节类型以及3字节的cts
//
// @return sps, pps: 复用传入参数`payload`的内存块
func ParseSpsPpsFromSeqHeaderWithoutMalloc(payload []byte) (sps, pps []byte, err error) {
	if len(payload) < 13 {
		return nil, nil, errors.New("buf too short")
	}
	if payload[0] != 0x17 || payload[1] != 0x00 || payload[2] != 0 || payload[3] != 0 || payload[4] != 0 {
		return nil, nil, errors.New("payload invalid")
	}

	index := 10
	numOfSps := int(payload[index] & 0x1F)
	index++
	if numOfSps != 1 {
		return nil, nil, errors.New("numOfSps != 1")
	}
	spsLength := int(binary.BigEndian.Uint16(payload[index:]))
	index += 2

	if len(payload) < 13+spsLength {
		return nil, nil, errors.New("buf too short")
	}

	sps = payload[index : index+spsLength]
	index += spsLength

	if len(payload) < 16+spsLength {
		return nil, nil, errors.New("buf too short")
	}

	numOfPps := int(payload[index] & 0x1F)
	index++
	if numOfPps != 1 {
		return nil, nil, errors.New("payload invalid")
	}
	ppsLength := int(binary.BigEndian.Uint16(payload[index:]))
	index += 2

	if len(payload) < 16+spsLength+ppsLength {
		return nil, nil, errors.New("buf too short")
	}

	pps = payload[index : index+ppsLength]
	return
}
